package com.da.tools;

import com.da.exception.OllamaClientException;
import com.da.result.StreamResult;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import java.util.Map;

/**
 * Ollama http请求工具
 *
 * @author da
 * @time 2024/7/3 下午 4:26
 */
public class Request {
    /**
     * 请求地址
     */
    private final String url;
    /**
     * 请求方法
     */
    private final String method;
    /**
     * post请求要携带的数据
     */
    private final Map<String, Object> data;
    /**
     * 响应状态码
     */
    private int code;

    /**
     * 构造函数
     *
     * @param url    请求地址
     * @param method 请求方法
     * @param data   请求携带的数据
     */
    public Request(String url, String method, Map<String, Object> data) {
        this.url = url;
        this.method = method;
        this.data = data;
    }

    /**
     * 构造函数
     *
     * @param url  请求地址
     * @param data 请求携带的数据
     */
    public Request(String url, Map<String, Object> data) {
        this(url, "POST", data);
    }

    /**
     * 发送请求得到响应结果字符串,解析失败就返回空字符串
     *
     * @return 响应结果字符串
     */
    public String send() {
        try {
            // 构建请求json
            String json = Parse.toJson(this.data);
            HttpURLConnection connection = (HttpURLConnection) new URL(this.url).openConnection();
            connection.setRequestMethod(this.method);
            if (this.method.equals("POST")) {
                connection.setRequestProperty("Content-Type", "application/json");
                connection.setDoOutput(true);
                OutputStream os = connection.getOutputStream();
                os.write(json.getBytes());
                os.flush();
                os.close();
            }
            this.code = connection.getResponseCode();
            StringBuilder sb = new StringBuilder();
            if (this.code == 200) {
                //非流式请求
                try (BufferedReader reader = new BufferedReader(new InputStreamReader(connection.getInputStream()))) {
                    sb.append(reader.readLine());
                } finally {
                    connection.disconnect();
                }
            }
            return sb.toString();
        } catch (IOException e) {
            throw new OllamaClientException(e);
        }
    }

    /**
     * 获取响应状态码
     *
     * @return 响应状态码
     */
    public int getCode() {
        return code;
    }

    /**
     * 发送请求得到响应结果流,解析失败就返回空流
     *
     * @return 响应结果流
     */
    public StreamResult stream(String type) {
        try {
            return new StreamResult(this.url, this.data, this.method, type);
        } catch (IOException e) {
            throw new OllamaClientException(e);
        }
    }
}
